#!/usr/bin/env python
from prometheus_msgs.msg import UAVState, UAVCommand, UAVSetup, UAVControlState
from enum import Enum

import rospy


class RunningState(Enum):
    WaitingForStart = 1
    TakeOff = 2
    Moving = 3
    Landing = 4


class MissionControl:

    def __init__(self, uav_id):
        self.node_name = "moving_node"
        rospy.init_node(self.node_name)
        self.prefix = "/uav" + str(uav_id)
        self.rate = rospy.Rate(20.0)

        self.state = UAVState()
        self.control_state = UAVControlState()

        self.command = UAVCommand()
        self.setup_msg = UAVSetup()

        self.state_sub = rospy.Subscriber(
            self.prefix + "/prometheus/state", UAVState, self.state_cb, queue_size=10
        )
        self.control_state_sub = rospy.Subscriber(
            self.prefix + "/prometheus/control_state",
            UAVControlState,
            self.control_state_cb,
            queue_size=10,
        )
        self.command_pub = rospy.Publisher(
            self.prefix + "/prometheus/command", UAVCommand, queue_size=10
        )
        self.setup_pub = rospy.Publisher(
            self.prefix + "/prometheus/setup", UAVSetup, queue_size=10
        )

        self.current_state = RunningState.WaitingForStart

        self.max_vel = None
        self.takeoff_height = None
        self.takeoff_vel = None

        self.moving_points = None
        self.current_index = 0
        self.vel_factor = None

        self.read_ros_params()
        self.print_params()

    def print_params(self):
        rospy.loginfo("Ros node init as {0}".format(self.node_name))
        rospy.loginfo("takeoff height is {0} [m]".format(self.takeoff_height))
        rospy.loginfo("Takeoff vel is {0} [m/s]".format(self.takeoff_vel))
        rospy.loginfo("Uav max vel is {0} [m/s]".format(self.max_vel))
        rospy.loginfo("Moving points is : {0}".format(self.moving_points))
        rospy.loginfo("UAV vel factor is {0}".format(self.vel_factor))

    def read_ros_params(self):

        self.moving_points = rospy.get_param("~moving_points", [[1, 1], [0, 0]])
        self.max_vel = rospy.get_param("~max_vel", 0.5)
        self.takeoff_vel = rospy.get_param("~takeoff_vel", 0.6)
        self.takeoff_height = rospy.get_param("~takeoff_height", 1.5)
        self.vel_factor = rospy.get_param("~vel_factor", 0.4)

    def state_cb(self, msg):
        self.state = msg

    def control_state_cb(self, msg):
        self.control_state = msg

    def get_vel(self, now_pos, target_pos):
        dis_vel = [
            (tar - now) * self.vel_factor for (tar, now) in zip(target_pos, now_pos)
        ]
        limit_vels = []

        for vel in dis_vel:
            if abs(vel) > self.max_vel:
                limit_vel = self.max_vel * vel / abs(vel)
            else:
                limit_vel = vel
            limit_vels.append(limit_vel)

        return limit_vels

    def run(self):
        while not rospy.is_shutdown():
            if self.current_state == RunningState.WaitingForStart:
                rospy.loginfo_once("Waiting for arm and COMMAND_CONTROL")
                if self.control_state.control_state == UAVControlState.COMMAND_CONTROL:
                    self.current_state = RunningState.TakeOff

            elif self.current_state == RunningState.TakeOff:
                rospy.loginfo_once("Enter Takeoff mode!")
                if abs(self.state.position[2] - self.takeoff_height) < 0.15:
                    self.current_state = RunningState.Moving
                else:
                    # vx, vy, vz = self.get_vel(self.state.position, [0, 0, self.takeoff_height])
                    self.command.Agent_CMD = UAVCommand.Move
                    self.command.Move_mode = UAVCommand.XYZ_VEL
                    self.command.velocity_ref[0] = 0
                    self.command.velocity_ref[1] = 0
                    self.command.velocity_ref[2] = self.takeoff_vel

            elif self.current_state == RunningState.Moving:
                rospy.loginfo_once("Enter Moving mode!")

                x, y = self.moving_points[self.current_index]
                dis2tar = abs(x - self.state.position[0]) + abs(
                    y - self.state.position[1]
                )
                if dis2tar < 0.2:
                    rospy.loginfo("Arrived at point {0} {1} [m]".format(x, y))
                    if self.current_index == len(self.moving_points) - 1:
                        self.current_state = RunningState.Landing
                    else:
                        self.current_index += 1
                else:
                    vx, vy, vz = self.get_vel(
                        self.state.position, [x, y, self.takeoff_height]
                    )
                    self.command.Agent_CMD = UAVCommand.Move
                    self.command.Move_mode = UAVCommand.XYZ_VEL
                    self.command.velocity_ref[0] = vx
                    self.command.velocity_ref[1] = vy
                    self.command.velocity_ref[2] = vz

            elif self.current_state == RunningState.Landing:
                rospy.loginfo_once("Enter Landing mode")

                self.command.Agent_CMD = UAVCommand.Land

            else:
                raise NotImplementedError

            if self.current_state == RunningState.Landing and not self.state.armed:
                rospy.loginfo("UAV in land state and is disarmed! Program exit!")
                break

            self.command.header.stamp = rospy.Time.now()
            self.command.Command_ID += 1
            self.command_pub.publish(self.command)

            # rospy.loginfo('vx: {0}, vy: {1}, vz: {2}'.format(self.command.velocity_ref[0], self.command.velocity_ref[1], self.command.velocity_ref[2]))

            self.rate.sleep()


if __name__ == "__main__":
    mission_control = MissionControl(1)
    mission_control.run()
